import numpy as np
try:
    import alphashape
except Exception:
    alphashape = None
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle
from matplotlib import animation
import torch
from shapely.geometry import Polygon, MultiPolygon


KINEMATIC_PARTICLE = 6
STATIONARY_PARTICLE = 3


def render_animation(
        positions: torch.tensor,
        particle_type,
        mpm_inputs,
        write_path,
        timestep_stride=5,
):

    # Start with computing displacement
    kinematic_positions, stationary_positions = get_positions_by_type(
        positions,
        particle_type)

    kinematic_positions = kinematic_positions.detach().cpu().numpy()
    stationary_positions = stationary_positions.detach().cpu().numpy()

    initial_kinematic_position = kinematic_positions[None, 0]
    initial_kinematic_positions = np.repeat(initial_kinematic_position, repeats=len(positions), axis=0)
    disp = np.linalg.norm(kinematic_positions - initial_kinematic_positions, axis=-1)

    fig = plt.figure()
    ax = fig.add_subplot(projection='3d')

    def animate(i):
        print(f"Render step {i}/{len(positions)}")
        fig.clear()

        cmap = plt.cm.viridis
        vmax = np.ndarray.flatten(disp).max()
        vmin = np.ndarray.flatten(disp).min()
        sampled_value = disp[i]

        # Note: z and y is interchanged to match taichi coordinate convention.
        ax = fig.add_subplot(projection='3d', autoscale_on=False)
        ax.set_xlim(mpm_inputs['sim_space'][0])
        ax.set_ylim(mpm_inputs['sim_space'][2])
        ax.set_zlim(mpm_inputs['sim_space'][1])
        ax.set_xlabel("x")
        ax.set_ylabel("z")
        ax.set_zlabel("y")
        ax.invert_zaxis()

        trj = ax.scatter(kinematic_positions[i][:, 0],
                         kinematic_positions[i][:, 2],
                         kinematic_positions[i][:, 1],
                         c=sampled_value, vmin=vmin, vmax=vmax, cmap=cmap, s=1)
        ax.scatter(stationary_positions[i][:, 0],
                   stationary_positions[i][:, 2],
                   stationary_positions[i][:, 1], c="black")
        fig.colorbar(trj)

        ax.set_box_aspect(
            aspect=(float(mpm_inputs['sim_space'][0][0]) - float(mpm_inputs['sim_space'][0][1]),
                    float(mpm_inputs['sim_space'][2][0]) - float(mpm_inputs['sim_space'][2][1]),
                    float(mpm_inputs['sim_space'][1][0]) - float(mpm_inputs['sim_space'][1][1])))
        ax.view_init(elev=20., azim=i*0.5)
        # ax.view_init(elev=20., azim=0.5)
        ax.grid(True, which='both')

    # Creat animation
    ani = animation.FuncAnimation(
        fig, animate, frames=np.arange(0, len(positions), timestep_stride), interval=20)

    ani.save(write_path, dpi=100, fps=30, writer='imagemagick')
    print(f"Animation saved to: {write_path}.gif")


def visualize_state(
        vis_data: dict,
        barrier_info: dict,
        mpm_inputs: dict,
        loss: float,
        write_path: str
):
    """
    Make runout comparison plot
    Args:
        vis_data (dict): data for visualization following a specific format
        barrier_info (dict): prescribed barrier geometry information.
            It is used to make a rectangular patch in fig to represent barriers
        mpm_inputs (dict): used to get simulation domain boundary
        loss (float):
        write_path (str):

    Returns:

    """

    barrier_height = barrier_info["barrier_height"]
    barrier_width = barrier_info["barrier_width"]

    vis_params = {
        "pred": {
            "label": "Pred",
            "perimeter_color": "black",
            "particle_color": "purple",
            "alpha": 0.5
        },
        "true": {
            "label": "True",
            "perimeter_color": "darkorange",
            "particle_color": "yellow",
            "alpha": 0.5
        }
    }

    # Init fig
    fig, ax = plt.subplots()

    for i, (key, value) in enumerate(vis_data.items()):
        # Preprocess data
        last_position = value["kinematic_positions"][-1].detach().cpu().numpy()
        runout_perimeter = None
        if alphashape is not None:
            try:
                runout_perimeter = alphashape.alphashape(
                    last_position[:, [0, 2]], alpha=20.0)
            except Exception:
                runout_perimeter = None
        if "runout_end" in value:
            runout_end = value["runout_end"].detach().cpu().numpy()
        if "centroid" in value:
            centroid = value["centroid"].detach().cpu().numpy()

        # Get rectangular patches for pred barriers
        barrier_patches = []
        barrier_locs = value["barrier_locs"].detach().cpu().numpy()
        for barrier_edge in barrier_locs:
            barrier_patch = Rectangle(
                xy=barrier_edge,
                width=barrier_width,
                height=barrier_width,
                edgecolor=vis_params[key]["perimeter_color"],
                fill=False,
                linewidth=1.5,
                zorder=10
            )
            barrier_patches.append(barrier_patch)

        # Plot runout
        ax.scatter(last_position[:, 0], last_position[:, 2],
                   alpha=0.3, s=2.0, c=vis_params[key]["particle_color"])

        # Plot runout end used in loss calculation
        if "runout_end" in value:
            ax.scatter(runout_end[:, 0], runout_end[:, 2],
                       s=5.0,
                       c=vis_params[key]["particle_color"],
                       edgecolors=vis_params[key]["perimeter_color"])
        if "centroid" in value:
            ax.scatter(centroid[0], centroid[2],
                       s=50.0,
                       c=vis_params[key]["particle_color"],
                       edgecolors=vis_params[key]["perimeter_color"],
                       zorder=11)

        ### Plot runout fill. Simple `PolygonPath` doesn't work in the current version of the package
        # ax.add_patch(PolygonPatch(runout_perimeter, alpha=0.2))
        if isinstance(runout_perimeter, Polygon):
            x, y = runout_perimeter.exterior.xy
            ax.fill(x, y, alpha=0.2,
                    color=vis_params[key]["perimeter_color"],
                    label=vis_params[key]["label"])
        elif isinstance(runout_perimeter, MultiPolygon):
            for geom in runout_perimeter.geoms:
                x, y = geom.exterior.xy
                ax.fill(x, y, alpha=0.2,
                        color=vis_params[key]["perimeter_color"],
                        label=vis_params[key]["label"])
        ###

        for barrier_patch in barrier_patches:
            ax.add_patch(barrier_patch)

    ax.set_xlim(mpm_inputs['sim_space'][0])
    ax.set_ylim(mpm_inputs['sim_space'][2])
    ax.set_aspect("equal")
    ax.set_title(f"Loss: {loss:.3e}")
    ax.legend()
    ax.grid(True)
    # plt.show()
    plt.savefig(write_path)


def get_features(
        kinematic_positions: torch.tensor,
        barrier_particles: torch.tensor,
        device: torch.device,
        n_init_steps: int = 6,
        sdf: torch.Tensor | None = None,
        sdf_path: str | None = None,
        sdf_key: str = "sdf",
        return_sdf: bool = False,
):
    """
    Prepare GNS rollout features.

    Args:
        kinematic_positions: (timesteps, nparticles, dims)
        barrier_particles: (n_stationary_particles, dims)
        device: torch device
        n_init_steps: number of initial steps to seed the simulator (default: 6)
        sdf: optional SDF tensor to be used as per-particle material property. Accepted shapes:
             - (nparticles, 8) or (nparticles, C)
             - (timesteps, nparticles, 8) or (timesteps, nparticles, C)
        sdf_path: optional path to load SDF from (.npz/.npy/.pt). If provided and `sdf` is None,
                  this function will attempt to load it.
        sdf_key: key name when loading from .npz (default: 'sdf').
        return_sdf: when True, return the processed per-particle scalar property as the 4th output.

    Returns:
        initial_positions: (nparticles_total, n_init_steps, dims)
        particle_type: (nparticles_total,)
        n_particles_per_example: (1,)
        material_property (optional): (nparticles_total,) – scalar derived from SDF for kinematic
                                      particles; zeros for stationary particles.
    """
    # Clamp n_init_steps to available history
    n_init_steps = int(max(1, min(n_init_steps, kinematic_positions.shape[0])))

    # Get initial positions
    initial_kinematic_positions = kinematic_positions[:n_init_steps].to(device)
    barrier_particles_unsqueeze = barrier_particles.unsqueeze(0)
    current_stationary_positions = barrier_particles_unsqueeze.repeat(n_init_steps, 1, 1)
    # Note that the shape is changed to (nparticles, timesteps, dims) from (timesteps, nparticles, dims)
    #   to follow GNS rollout function convention
    initial_positions = torch.concat(
        (initial_kinematic_positions, current_stationary_positions), dim=1
    ).permute(1, 0, 2).to(torch.float32).contiguous()

    # Make particle type
    kinematic_particle_type = torch.full(
        (initial_kinematic_positions.shape[1], ), KINEMATIC_PARTICLE)
    stationary_particle_type = torch.full(
        (current_stationary_positions.shape[1], ), STATIONARY_PARTICLE)
    particle_type = torch.cat(
        (kinematic_particle_type, stationary_particle_type)).contiguous().to(device)

    # Get number of particles
    n_particles_per_example = torch.tensor([len(particle_type)], dtype=torch.int32).to(device)

    # 创建8维的material_property (SDF场)
    n_particles = len(particle_type)
    material_property = torch.zeros((n_particles, 8), device=device)

    # 为运动粒子设置SDF值
    # 1-3维: 外边界的梯度向量分量
    # 4维: 外边界的SDF值
    # 5-7维: 内边界的梯度向量分量
    # 8维: 内边界的SDF值
    kinematic_count = len(kinematic_particle_type)
    material_property[:kinematic_count, :3] = torch.tensor([0.0, 0.0, 0.0], device=device)  # 外边界梯度
    material_property[:kinematic_count, 3] = -1.0  # 外边界SDF值
    material_property[:kinematic_count, 4:7] = torch.tensor([0.0, 0.0, 0.0], device=device)  # 内边界梯度
    material_property[:kinematic_count, 7] = 1.0  # 内边界SDF值

    # 为障碍物粒子设置SDF值
    stationary_count = len(stationary_particle_type)
    material_property[stationary_count:, :3] = torch.tensor([0.0, 0.0, 0.0], device=device)  # 外边界梯度
    material_property[stationary_count:, 3] = 1.0  # 外边界SDF值
    material_property[stationary_count:, 4:7] = torch.tensor([0.0, 0.0, 0.0], device=device)  # 内边界梯度
    material_property[stationary_count:, 7] = -1.0  # 内边界SDF值

    # # Optionally load/process SDF as material property
    # material_property = None
    # if sdf is None and sdf_path is not None:
    #     try:
    #         if sdf_path.endswith(".npz"):
    #             data = np.load(sdf_path, allow_pickle=True)
    #             if sdf_key in data:
    #                 sdf_np = data[sdf_key]
    #             else:
    #                 # fallback to first array in npz
    #                 sdf_np = list(data.items())[0][1]
    #         elif sdf_path.endswith(".npy"):
    #             sdf_np = np.load(sdf_path, allow_pickle=True)
    #         elif sdf_path.endswith(".pt") or sdf_path.endswith(".pth"):
    #             sdf_np = torch.load(sdf_path)
    #             if isinstance(sdf_np, torch.Tensor):
    #                 sdf = sdf_np
    #                 sdf_np = None
    #         else:
    #             sdf_np = None
    #
    #         if sdf is None and sdf_np is not None:
    #             sdf = torch.tensor(sdf_np)
    #     except Exception:
    #         sdf = None
    #
    # if sdf is not None:
    #     # Normalize SDF shapes to (n_init_steps, n_kinematic, C)
    #     if sdf.dim() == 2:  # (nparticles, C)
    #         # assume kinematic particles count
    #         sdf = sdf.unsqueeze(0).repeat(n_init_steps, 1, 1)
    #     elif sdf.dim() == 3:  # (timesteps, nparticles, C)
    #         sdf = sdf[:n_init_steps]
    #     else:
    #         sdf = None
    #
    #     if sdf is not None:
    #         sdf = sdf.to(dtype=torch.float32, device=device)
    #         n_kinematic = initial_kinematic_positions.shape[1]
    #         # Reduce channel dimension to scalar per particle via L2 norm
    #         sdf_scalar = torch.linalg.norm(sdf, dim=-1)  # (n_init_steps, n_kinematic)
    #         # Use last step scalar as property for kinematic particles
    #         kinematic_property = sdf_scalar[-1]  # (n_kinematic,)
    #         # For stationary particles, set zeros
    #         stationary_property = torch.zeros(current_stationary_positions.shape[1], device=device, dtype=torch.float32)
    #         material_property = torch.cat([kinematic_property, stationary_property], dim=0)
    #
    # if return_sdf:
    #     return initial_positions, particle_type, n_particles_per_example, material_property
    # else:
    return initial_positions, particle_type, n_particles_per_example


def render_mixing_gif_yoz(
        positions: torch.tensor,
        particle_type: torch.tensor,
        write_path: str,
        timestep_stride: int = 1,
        type1: int = 6,
        type2: int = 0,
        figsize=(6, 6),
        dpi: int = 120,
        # Overlays (optional)
        cylinder_center=(0.055, 0.055),
        inner_radius: float | None = 0.05,
        obstacle_center: tuple[float, float] | None = None,
        obstacle_len_y: float | None = None,
        obstacle_len_z: float | None = None,
        obstacle_angle_rad: float | None = None,
        # Styles (optional)
        ring_color='cyan',
        ring_linewidth: float = 2.0,
        ring_alpha: float = 0.8,
        rect_color='red',
        rect_linewidth: float = 2.5,  # 增加线宽使其更明显
        rect_alpha: float = 1.0,  # 增加不透明度
):
    """Render YOZ plane mixing animation for two particle types.

    Args:
        positions: (timesteps, nparticles, 3)
        particle_type: (nparticles,)
        write_path: output path without extension
        timestep_stride: frame stride
        type1: first particle type id (e.g., 6)
        type2: second particle type id (e.g., 0)
    """
    pos = positions.detach().cpu().numpy()
    ptype = particle_type.detach().cpu().numpy()

    idx_type1 = np.where(ptype == type1)[0]
    idx_type2 = np.where(ptype == type2)[0]
    idx_type3 = np.where(ptype == 3)[0]  # 障碍/静止粒子

    fig, ax = plt.subplots(figsize=figsize)

    # Determine yz bounds
    y_all = pos[:, :, 1]
    z_all = pos[:, :, 2]
    y_min, y_max = float(np.min(y_all)), float(np.max(y_all))
    z_min, z_max = float(np.min(z_all)), float(np.max(z_all))

    # Auto-estimate ring radius if not provided
    cy_auto, cz_auto = cylinder_center
    if inner_radius is None:
        rr = np.sqrt((y_all - cy_auto) ** 2 + (z_all - cz_auto) ** 2)
        ring_radius = float(np.max(rr)) * 1.01
    else:
        ring_radius = float(inner_radius)

    def animate(i):
        ax.clear()
        t = i
        yy = pos[t, :, 1]
        zz = pos[t, :, 2]
        ax.scatter(yy[idx_type1], zz[idx_type1], s=2, c='tab:orange', label=f'Type {type1}', alpha=0.7)
        ax.scatter(yy[idx_type2], zz[idx_type2], s=4, c='tab:blue', label=f'Type {type2}', alpha=0.7)
        if len(idx_type3) > 0:
            ax.scatter(yy[idx_type3], zz[idx_type3], s=2, c='k', label='Type 3', alpha=0.7)

        # Outer cylinder ring
        if ring_radius is not None and np.isfinite(ring_radius) and ring_radius > 0:
            cy, cz = cylinder_center
            ring = plt.Circle((cy, cz), ring_radius, color=ring_color, fill=False, linewidth=ring_linewidth,
                              alpha=ring_alpha)
            ax.add_patch(ring)

            # 添加圆筒中心标记
            ax.scatter([cy], [cz], s=20, c=ring_color, marker='x')

        # Obstacle rectangle overlay (rotated in YOZ)
        if (obstacle_center is not None and
                obstacle_len_y is not None and
                obstacle_len_z is not None and
                obstacle_angle_rad is not None):
            cyy, czz = obstacle_center
            hy = float(obstacle_len_y) / 2.0
            hz = float(obstacle_len_z) / 2.0
            phi = np.pi/2 - float(obstacle_angle_rad)

            # 创建旋转矩阵
            cos_a = np.cos(phi)
            sin_a = np.sin(phi)

            # 定义矩形的四个角点（相对于中心）
            corners_local = np.array([
                [hy, hz],
                [-hy, hz],
                [-hy, -hz],
                [hy, -hz],
                [hy, hz]  # 闭合矩形
            ])

            # 应用旋转和平移
            rect_y, rect_z = [], []
            for corner in corners_local:
                y_r, z_r = corner
                y_w = y_r * cos_a - z_r * sin_a + cyy
                z_w = y_r * sin_a + z_r * cos_a + czz
                rect_y.append(y_w)
                rect_z.append(z_w)

            # 绘制矩形边界框
            ax.plot(rect_y, rect_z, color=rect_color, linewidth=rect_linewidth, alpha=rect_alpha)

            # 添加填充区域以更清楚地显示无颗粒区域
            from matplotlib.patches import Polygon
            polygon = Polygon(list(zip(rect_y[:-1], rect_z[:-1])),
                              facecolor='none',
                              edgecolor=rect_color,
                              linewidth=rect_linewidth,
                              alpha=rect_alpha)
            ax.add_patch(polygon)

        # 设置坐标范围
        buffer = ring_radius * 0.1
        cy, cz = cylinder_center
        ax.set_xlim(cy - ring_radius - buffer, cy + ring_radius + buffer)
        ax.set_ylim(cz - ring_radius - buffer, cz + ring_radius + buffer)

        ax.set_xlabel('Y')
        ax.set_ylabel('Z')
        ax.set_title(f'YOZ Mixing (t={t})')
        ax.legend(loc='upper right', markerscale=3)
        ax.set_aspect('equal', adjustable='box')

    frames = list(range(0, pos.shape[0], max(1, timestep_stride)))
    ani = animation.FuncAnimation(fig, animate, frames=frames, interval=50)
    # Try ImageMagick first, fallback to Pillow
    try:
        ani.save(f'{write_path}.gif', dpi=dpi, fps=30, writer='imagemagick')
    except Exception:
        ani.save(f'{write_path}.gif', dpi=dpi, fps=30, writer='pillow')
    plt.close(fig)
    print(f"Saved YOZ mixing GIF to: {write_path}.gif")


def render_yoz_png(
        positions: torch.tensor,
        particle_type: torch.tensor,
        t: int,
        write_path: str,
        cylinder_center=(0.055, 0.055),
        inner_radius: float | None = None,
        obstacle_center: tuple[float, float] | None = None,
        obstacle_len_y: float | None = None,
        obstacle_len_z: float | None = None,
        obstacle_angle_rad: float | None = None,
        figsize=(5, 5),
        dpi: int = 150,
        type1: int = 6,
        type2: int = 0,
        # Styles 
        ring_color='cyan',
        ring_linewidth: float = 2.0,
        ring_alpha: float = 0.8,
        rect_color='red',
        rect_linewidth: float = 1.5,
        rect_alpha: float = 0.9,
):
    """Render one YOZ snapshot PNG with per-type colors and boundary overlays.

    Args:
        positions: (timesteps, nparticles, 3)
        particle_type: (nparticles,)
        t: timestep index to render
        write_path: output png path
        cylinder_center: (cy, cz)
        inner_radius: inner cylinder radius; when provided draw a ring
        obstacle_center: (oy, oz)
        obstacle_len_y: rectangle size along Y
        obstacle_len_z: rectangle size along Z
        obstacle_angle_rad: angle around +X; when provided draw rotated rectangle
    """
    pos = positions.detach().cpu().numpy()
    ptype = particle_type.detach().cpu().numpy()

    t = int(max(0, min(t, pos.shape[0] - 1)))
    yy = pos[t, :, 1]
    zz = pos[t, :, 2]

    idx1 = np.where(ptype == type1)[0]
    idx2 = np.where(ptype == type2)[0]
    idx3 = np.where(ptype == 3)[0]

    fig, ax = plt.subplots(figsize=figsize)
    ax.scatter(yy[idx1], zz[idx1], s=2, c='tab:orange', alpha=0.7, label=f'Type {type1}')
    ax.scatter(yy[idx2], zz[idx2], s=2, c='tab:blue', alpha=0.7, label=f'Type {type2}')
    if len(idx3) > 0:
        ax.scatter(yy[idx3], zz[idx3], s=2, c='k', alpha=0.7, label='Type 3')

    # Outer cylinder ring (auto-estimate if missing)
    cy, cz = cylinder_center
    if inner_radius is None:
        rr = np.sqrt((yy - cy) ** 2 + (zz - cz) ** 2)
        ring_r = float(np.max(rr)) * 1.01
    else:
        ring_r = float(inner_radius)
    if np.isfinite(ring_r) and ring_r > 0:
        ring = plt.Circle((cy, cz), ring_r, color=ring_color, fill=False, linewidth=ring_linewidth, alpha=ring_alpha)
        ax.add_patch(ring)

    # Obstacle rectangle overlay
    if obstacle_center is not None and obstacle_len_y is not None and obstacle_len_z is not None and obstacle_angle_rad is not None:
        cyy, czz = obstacle_center
        hy = float(obstacle_len_y) / 2.0
        hz = float(obstacle_len_z) / 2.0
        # angle around +X to a YOZ-aligned rectangle needs conversion used elsewhere
        phi = (np.pi / 2.0) - float(obstacle_angle_rad)
        cos_a = np.cos(phi)
        sin_a = np.sin(phi)
        corners_local = [
            (+hy, +hz),
            (-hy, +hz),
            (-hy, -hz),
            (+hy, -hz),
            (+hy, +hz),
        ]
        rect_y, rect_z = [], []
        for y_r, z_r in corners_local:
            y_w = y_r * cos_a - z_r * sin_a + cyy
            z_w = y_r * sin_a + z_r * cos_a + czz
            rect_y.append(y_w)
            rect_z.append(z_w)
        ax.plot(rect_y, rect_z, color=rect_color, linewidth=rect_linewidth, alpha=rect_alpha)

    ax.set_aspect('equal', adjustable='box')
    ax.set_xlabel('Y')
    ax.set_ylabel('Z')
    ax.set_title(f'YOZ Mixing (t={t})')
    ax.legend(loc='upper right', markerscale=3)
    fig.tight_layout()
    fig.savefig(write_path, dpi=dpi)
    plt.close(fig)

def fill_rectangle_yz_with_particles(
        len_x: float,
        len_y: float,
        len_z: float,
        spacing: float,
        center_x: float = 0.0,
        center_y: float = 0.0,
        center_z: float = 0.0,
        random_disturbance_mag: float = 0.0
):
    """Generate a YOZ-rectangle as stationary particle cloud at fixed x.

    Returns:
        torch.tensor of shape (nparticles, 3)
    """
    n_particles_x = int(max(1, len_x / spacing))
    n_particles_y = int(max(1, len_y / spacing))
    n_particles_z = int(max(1, len_z / spacing))

    x0 = center_x - len_x / 2.0
    y0 = center_y - len_y / 2.0
    z0 = center_z - len_z / 2.0
    x_coords = np.linspace(x0 + spacing / 2, x0 + len_x - spacing / 2, n_particles_x)
    y_coords = np.linspace(y0 + spacing / 2, y0 + len_y - spacing / 2, n_particles_y)
    z_coords = np.linspace(z0 + spacing / 2, z0 + len_z - spacing / 2, n_particles_z)

    X, Y, Z = np.meshgrid(x_coords, y_coords, z_coords, indexing='ij')
    particles = np.stack([
        X.flatten(),
        Y.flatten(),
        Z.flatten()
    ], axis=1)

    if random_disturbance_mag > 0:
        disturbance = np.random.uniform(
            -spacing/2 * random_disturbance_mag,
            spacing/2 * random_disturbance_mag,
            size=particles.shape
        )
        particles = particles + disturbance

    return torch.tensor(particles, dtype=torch.float32)


# def fill_cuboid_with_particles(
#         len_x: float,
#         len_y: float,
#         len_z: float,
#         spacing: float,
#         center_x: float = 0.015,
#         center_y: float = 0.055,
#         center_z: float = 0.055,
#         random_disturbance_mag: float = 0.0
# ) -> torch.Tensor:
#     """Generate a 3D cuboid particle cloud with specified dimensions.
#
#     Args:
#         len_x: Length in X-dimension (0-0.03 in your case)
#         len_y: Length in Y-dimension
#         len_z: Length in Z-dimension
#         spacing: Distance between particles
#         center_x: Center coordinate of cuboid in X
#         center_y: Center coordinate of cuboid in Y
#         center_z: Center coordinate of cuboid in Z
#         random_disturbance_mag: Magnitude of random position disturbance
#
#     Returns:
#         torch.Tensor of shape (n_particles, 3)
#     """
#     # Calculate number of particles in each dimension
#     n_x = max(1, int(len_x / spacing))
#     n_y = max(1, int(len_y / spacing))
#     n_z = max(1, int(len_z / spacing))
#
#     # Create coordinate grids for each dimension
#     x0 = center_x - len_x / 2.0
#     y0 = center_y - len_y / 2.0
#     z0 = center_z - len_z / 2.0
#
#     x_coords = np.linspace(x0 + spacing / 2, x0 + len_x - spacing / 2, n_x)
#     y_coords = np.linspace(y0 + spacing / 2, y0 + len_y - spacing / 2, n_y)
#     z_coords = np.linspace(z0 + spacing / 2, z0 + len_z - spacing / 2, n_z)
#
#     # Create 3D grid of positions
#     X, Y, Z = np.meshgrid(x_coords, y_coords, z_coords, indexing='ij')
#
#     # Flatten and stack into [x, y, z] columns
#     particles = np.stack([
#         X.flatten(),
#         Y.flatten(),
#         Z.flatten()
#     ], axis=1)
#
#     # Apply random position disturbance if needed
#     if random_disturbance_mag > 0:
#         disturbance = np.random.uniform(
#             -spacing / 2 * random_disturbance_mag,
#             spacing / 2 * random_disturbance_mag,
#             size=particles.shape
#         )
#         particles += disturbance
#
#     return torch.tensor(particles, dtype=torch.float32)


def locate_barrier_particles(
        barrier_particles: torch.tensor,
        barrier_locations: torch.tensor,
        base_height: torch.tensor
):
    # Make current barrier particles with the current locations
    current_barrier_groups = []
    for loc in barrier_locations:
        moved_particles = barrier_particles.clone()
        # move `barrier_particles` corresponding to the current location
        moved_particles[:, 0] += loc[0]
        moved_particles[:, 1] += base_height
        moved_particles[:, 2] += loc[1]
        current_barrier_groups.append(moved_particles)
    current_barrier_particles = torch.vstack(current_barrier_groups)

    return current_barrier_particles


def get_runout_end(last_kinematic_positions, n_farthest_particles):
    """
    Get particle positions at the last timestep for `n_farthest_particles` to compute loss
    Args:
        last_kinematic_positions (torch.tensor): particle positions at the last timestep of flow
        n_farthest_particles (int): n particles to sample from the farthest x positions

    Returns:
    `n_farthest_particles` sampled from the last particle positions
    """

    # sort particle positions at the last timestep in ascending order
    kinematic_particles_idx_sorted = torch.argsort(last_kinematic_positions[:, 0])
    # get idx of the last `n_farthest_particles`
    runout_end_idx = kinematic_particles_idx_sorted[-n_farthest_particles:]
    # get `runout_end` consisting of `n_farthest_particles` of the runout
    runout_end = last_kinematic_positions[runout_end_idx, :]

    return runout_end


def get_positions_by_type(
        positions: torch.tensor,
        particle_types: torch.tensor):

    kinematic_particles_idx = torch.where(particle_types == KINEMATIC_PARTICLE)[0]
    stationary_particles_idx = torch.where(particle_types == STATIONARY_PARTICLE)[0]
    kinematic_positions = positions[:, kinematic_particles_idx, :]
    stationary_positions = positions[:, stationary_particles_idx, :]

    return kinematic_positions, stationary_positions


def fill_cuboid_with_particles(
        len_x,
        len_y,
        len_z,
        spacing,
        random_disturbance_mag=0.8,
):
    """
    Fill a cuboid with particles using PyTorch tensors and linspace to avoid graph breakage.

    Parameters:
    - len_x: Tensor scalar representing the length of the cuboid along the x-axis.
    - len_y: Tensor scalar representing the height of the cuboid along the y-axis.
    - len_z: Tensor scalar representing the width of the cuboid along the z-axis.
    - spacing: Tensor scalar representing the distance between adjacent particles.
    - random_disturbance_mag: Tensor scalar for randomization magnitude to disturb particle distribution

    Returns:
    - A PyTorch tensor, where each row represents the (x, y, z) coordinates of a particle.
    """

    # Compute nparticles per dim
    n_particles_x = int(len_x / spacing)
    n_particles_y = int(len_y / spacing)
    n_particles_z = int(len_z / spacing)

    # Generate grid of points
    # Assuming particles are located as follows: | -- * ---- * -- |,
    #   cell size is cellsize=domain/resolution (e.g., 2.0/64),
    #   particle spacing is spacing=2.0/64/2,
    #   starting point of particle location is spacing/2
    x = np.linspace(spacing/2, len_x - spacing/2, n_particles_x)
    y = np.linspace(spacing/2, len_y - spacing/2, n_particles_y)
    z = np.linspace(spacing/2, len_z - spacing/2, n_particles_z)

    # Generate meshgrid
    X, Y, Z = np.meshgrid(x, y, z, indexing='ij')

    # Flatten and stack to create a list of coordinates
    particles = np.stack([X.flatten(), Y.flatten(), Z.flatten()], axis=1)

    # Disturbance
    disturbance = np.random.uniform(-spacing/2 * random_disturbance_mag,
                                    spacing/2 * random_disturbance_mag,
                                    size=particles.shape)
    disturbed_particles = particles + disturbance
    torch_particles = torch.tensor(disturbed_particles, requires_grad=True)

    return torch_particles


def visualize_particles(particles_array):
    """
    Visualize the particles within a cuboid using a 3D scatter plot.

    Parameters:
    - particles_array: A NumPy array with shape (N, 3) representing the particles' coordinates.
    """
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')

    # Scatter plot of particles
    ax.scatter(particles_array[:, 0], particles_array[:, 1], particles_array[:, 2])

    # Labeling the axes
    ax.set_xlabel('X axis')
    ax.set_ylabel('Y axis')
    ax.set_zlabel('Z axis')

    # Set equal aspect ratio for all axes
    max_range = np.array([particles_array[:, 0].max() - particles_array[:, 0].min(),
                          particles_array[:, 1].max() - particles_array[:, 1].min(),
                          particles_array[:, 2].max() - particles_array[:, 2].min()]).max() / 2.0

    mid_x = (particles_array[:, 0].max() + particles_array[:, 0].min()) * 0.5
    mid_y = (particles_array[:, 1].max() + particles_array[:, 1].min()) * 0.5
    mid_z = (particles_array[:, 2].max() + particles_array[:, 2].min()) * 0.5

    ax.set_xlim(mid_x - max_range, mid_x + max_range)
    ax.set_ylim(mid_y - max_range, mid_y + max_range)
    ax.set_zlim(mid_z - max_range, mid_z + max_range)

    plt.show()

# # Example usage
# # Example usage
# len_x = 0.1 # Length of the cuboid along the x-axis
# len_y = 0.3   # Height of the cuboid along the y-axis
# len_z = 0.1   # Width of the cuboid along the z-axis
# spacing = 2.0/64/2  # Spacing between particles
# random_disturbance_mag = 0.8  # Random disturbance magnitude
#
# # Generate the particles
# particles = fill_cuboid_with_particles(len_x, len_y, len_z, spacing)
# visualize_particles(particles.detach().numpy())
# a=1

